import pandas as pd
import numpy as np
import matplotlib
# 使用非交互式后端以加快执行速度
matplotlib.use('Agg')
import matplotlib.pyplot as plt
# 提前设置matplotlib全局字体配置，确保中文字体正常显示
# 首先尝试使用Windows系统中肯定存在的字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'Microsoft YaHei', 'SimSun']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 强制设置matplotlib的字体缓存
import matplotlib.font_manager as fm
# 清除字体缓存
fm._get_font.cache_clear()
import seaborn as sns
import mplfinance as mpf
import os
import sys

# 设置中文字体
sns.set(font='SimHei', font_scale=1.2)
plt.rcParams['font.sans-serif'] = ['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC']
plt.rcParams['axes.unicode_minus'] = False

# 获取当前脚本所在目录
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 1. 读取文件，观察数据形状，输出数据前5行，并对定量变量进行描述性分析
def task1(file_path):
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path)
        
        print("=== 任务1：数据基本信息和描述性分析 ===")
        print(f"数据形状：{df.shape}")
        print("\n数据前5行：")
        print(df.head())
        
        # 查看数据类型
        print("\n数据类型：")
        print(df.dtypes)
        
        # 对定量变量进行描述性分析
        numeric_cols = df.select_dtypes(include=['float64', 'int64']).columns
        print("\n定量变量的描述性统计：")
        print(df[numeric_cols].describe())
        
        return df
    except Exception as e:
        print(f"任务1执行出错: {str(e)}")
        return None

# 2. 计算变量之间的相关系数，并绘制相关系数热力图
def task2(df):
    try:
        print("\n=== 任务2：相关系数分析 ===")
        
        # 选择定量变量计算相关系数
        numeric_cols = df.select_dtypes(include=['float64', 'int64']).columns
        corr_matrix = df[numeric_cols].corr()
        
        print("相关系数矩阵：")
        print(corr_matrix)
        
        # 绘制热力图
        plt.figure(figsize=(12, 8))
        sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', square=True, linewidths=.5, fmt='.2f')
        plt.title('比亚迪股票各变量之间的相关系数热力图')
        plt.tight_layout()
        
        # 保存到当前目录
        plt.savefig(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'correlation_heatmap.png'), dpi=300)
        print("相关系数热力图已保存")
        plt.close()
    except Exception as e:
        print(f"任务2执行出错: {str(e)}")

# 3. 绘制2023年全年的收盘价时序图
def task3(df):
    try:
        print("\n=== 任务3：2023年全年收盘价时序图 ====")
        
        # 确保日期列是datetime类型
        date_column = None
        # 尝试不同的可能的日期列名
        for col in ['日期', '交易日期', '时间']:
            if col in df.columns:
                date_column = col
                break
        
        if date_column:
            df[date_column] = pd.to_datetime(df[date_column])
            # 筛选2023年的数据
            df_2023 = df[(df[date_column].dt.year == 2023)]
            
            # 获取收盘价列
            close_column = None
            for col in ['收盘价', '收盘', 'close', 'Close']:
                if col in df.columns:
                    close_column = col
                    break
            
            if close_column:
                # 绘制收盘价时序图
                plt.figure(figsize=(15, 8))
                plt.plot(df_2023[date_column], df_2023[close_column], marker='o', markersize=3, linestyle='-', color='blue')
                plt.title('比亚迪股票2023年全年收盘价时序图')
                plt.xlabel('日期')
                plt.ylabel('收盘价（元）')
                plt.grid(True, alpha=0.3)
                plt.tight_layout()
                
                # 保存到当前目录
                plt.savefig(os.path.join(os.path.dirname(os.path.abspath(__file__)), '2023_close_price.png'), dpi=300)
                print("2023年全年收盘价时序图已保存")
                plt.close()
            else:
                print(f"未找到收盘价列，可用列名: {list(df.columns)}")
        else:
            print("未找到日期列")
    except Exception as e:
        print(f"任务3执行出错: {str(e)}")
        import traceback
        traceback.print_exc()

# 4. 绘制2023年4月到2023年6月的K线图，并进行分析
def task4(df):
    try:
        print("\n=== 任务4：2023年4月-6月K线图分析 ===")
        
        # 确保日期列是datetime类型
        date_column = None
        for col in ['日期', '交易日期', '时间']:
            if col in df.columns:
                date_column = col
                break
        
        if date_column:
            df[date_column] = pd.to_datetime(df[date_column])
            # 筛选2023年4月到6月的数据
            start_date = '2023-04-01'
            end_date = '2023-06-30'
            df_kline = df[(df[date_column] >= start_date) & (df[date_column] <= end_date)].copy()
            
            # 检查是否有足够的数据
            if len(df_kline) == 0:
                print(f"没有找到{start_date}到{end_date}期间的数据")
                return
            
            # 重命名列以适配mplfinance，使用实际列名
            column_mapping = {}
            for key, value in {'日期': 'Date', '开盘价': 'Open', '最高价': 'High', '最低价': 'Low', '收盘价': 'Close', '成交量': 'Volume'}.items():
                for col in df_kline.columns:
                    if key in col or key.lower() in col.lower():
                        column_mapping[col] = value
                        break
            
            # 如果找不到所有必要的列，尝试使用通用方法
            if not all(x in column_mapping.values() for x in ['Open', 'High', 'Low', 'Close']):
                # 假设第1列是日期，第2-5列是OHLC
                numeric_cols = df_kline.select_dtypes(include=['float64', 'int64']).columns[:4]
                if len(numeric_cols) >= 4:
                    column_mapping = {
                        date_column: 'Date',
                        numeric_cols[0]: 'Open',
                        numeric_cols[1]: 'High',
                        numeric_cols[2]: 'Low',
                        numeric_cols[3]: 'Close'
                    }
                    # 查找成交量列
                    for col in df_kline.columns:
                        if '成交量' in col or 'volume' in col.lower():
                            column_mapping[col] = 'Volume'
                            break
            
            # 重命名列
            df_kline_renamed = df_kline.rename(columns=column_mapping)
            
            # 确保必要的列存在
            required_cols = ['Date', 'Open', 'High', 'Low', 'Close']
            missing_cols = [col for col in required_cols if col not in df_kline_renamed.columns]
            
            if missing_cols:
                print(f"缺少必要的列：{missing_cols}")
                print("可用列：", df_kline_renamed.columns.tolist())
                return
            
            df_kline_renamed = df_kline_renamed.set_index('Date')
            
            # 确保数据类型正确
            for col in ['Open', 'High', 'Low', 'Close']:
                df_kline_renamed[col] = pd.to_numeric(df_kline_renamed[col], errors='coerce')
            
            # 创建一个完全自定义的mplfinance样式，明确指定中文字体
            # 修复了marketcolors缺少alpha参数的问题
            custom_style = {
                'base_mpl_style': 'default',
                'marketcolors': {
                    'candle': {'up': 'red', 'down': 'green'},
                    'edge': {'up': 'red', 'down': 'green'},
                    'wick': {'up': 'red', 'down': 'green'},
                    'volume': {'up': 'red', 'down': 'green'},
                    'ohlc': {'up': 'red', 'down': 'green'},
                    'vcedge': {'up': 'red', 'down': 'green'},
                    'alpha': 1.0,
                    'vcdopcod': 'volume'  # 添加成交量颜色决定参数
                },
                'mavcolors': ['blue', 'orange', 'purple'],
                'facecolor': 'white',
                'edgecolor': 'gray',
                'figcolor': 'white',
                'gridcolor': 'lightgray',
                'gridstyle': '--',
                'y_on_right': False,
                'rc': {
                    'font.family': ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS'],
                    'axes.unicode_minus': False
                }
            }
            
            # 使用自定义样式绘制K线图
            mpf.plot(df_kline_renamed, 
                     type='candle', 
                     volume='Volume' in df_kline_renamed.columns, 
                     title='比亚迪2023年4-6月K线', 
                     ylabel='价格(元)', 
                     ylabel_lower='成交量',
                     style=custom_style,
                     savefig=dict(fname=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'kline_chart.png'), dpi=300))
            print("2023年4月-6月K线图已保存")
            
            # K线图分析
            print("\nK线图分析：")
            # 计算基本统计信息
            max_price = df_kline_renamed['High'].max()
            min_price = df_kline_renamed['Low'].min()
            avg_price = df_kline_renamed['Close'].mean()
            price_change = df_kline_renamed['Close'].iloc[-1] - df_kline_renamed['Open'].iloc[0]
            percent_change = (price_change / df_kline_renamed['Open'].iloc[0]) * 100
            
            print(f"- 价格区间：{min_price:.2f} - {max_price:.2f}元")
            print(f"- 平均收盘价：{avg_price:.2f}元")
            print(f"- 区间涨跌：{price_change:.2f}元 ({percent_change:.2f}%)")
            
            # 分析成交量变化
            if 'Volume' in df_kline_renamed.columns:
                avg_volume = df_kline_renamed['Volume'].mean()
                max_volume = df_kline_renamed['Volume'].max()
                print(f"- 平均成交量：{avg_volume:.0f}")
                print(f"- 最大成交量：{max_volume:.0f}")
            
            # 简单趋势分析
            if percent_change > 5:
                print("- 趋势分析：该时间段内股价呈明显上涨趋势")
            elif percent_change < -5:
                print("- 趋势分析：该时间段内股价呈明显下跌趋势")
            else:
                print("- 趋势分析：该时间段内股价整体呈震荡走势")
        else:
            print("未找到日期列")
    except Exception as e:
        print(f"任务4执行出错: {str(e)}")
        import traceback
        traceback.print_exc()

# 主函数
def main():
    try:
        # 获取当前目录
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 文件路径 - 使用相对路径
        file_path = os.path.join(current_dir, '比亚迪后复权历史行情数据.xlsx')
        
        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"错误：找不到文件 {file_path}")
            print(f"当前工作目录: {os.getcwd()}")
            print(f"当前脚本目录: {current_dir}")
            return
        
        print("=== 开始执行比亚迪股票价格走势分析与可视化任务 ==")
        
        # 执行任务1
        df = task1(file_path)
        
        if df is not None and not df.empty:
            # 检查DataFrame列名
            print("\n=== 数据列名信息 ===")
            print(f"数据列名：{list(df.columns)}")
            
            # 执行任务2
            task2(df)
            
            # 执行任务3
            task3(df)
            
            # 执行任务4
            task4(df)
        else:
            print("未能成功加载数据，无法继续执行后续任务")
        
        print("\n=== 所有任务执行完毕 ===")
        
        # 列出生成的文件
        print("\n生成的文件：")
        for file in ['correlation_heatmap.png', '2023_close_price.png', 'kline_chart.png']:
            file_path = os.path.join(current_dir, file)
            if os.path.exists(file_path):
                print(f"- {file}")
    except Exception as e:
        print(f"主程序执行出错: {str(e)}")

if __name__ == "__main__":
    main()